#include "LeNet.h"
#include "ParamsLoader.h"
#include "BasicOp.h"
#include "AdvanceOp.h"

const std::string lables[] = { "plane", "car", "bird", "cat", "deer","dog", "frog", "horse", "ship", "truck" };

LeNet::LeNet(std::string paramsFile)
{
	ParamsLoader paramsLoader;
	paramsLoader.LoadFromFile(paramsFile);

	this->conv1Weight = &paramsLoader.LoadTensor4DFloat("conv1.weight");
	this->conv1Bias = &paramsLoader.LoadTensor1DFloat("conv1.bias");
	this->conv2Weight = &paramsLoader.LoadTensor4DFloat("conv2.weight");
	this->conv2Bias = &paramsLoader.LoadTensor1DFloat("conv2.bias");
	this->linear1Weight = &paramsLoader.LoadTensor2DFloat("fc1.weight");
	this->linear1Bias = &paramsLoader.LoadTensor1DFloat("fc1.bias");
	this->linear2Weight = &paramsLoader.LoadTensor2DFloat("fc2.weight");
	this->linear2Bias = &paramsLoader.LoadTensor1DFloat("fc2.bias");
	this->linear3Weight = &paramsLoader.LoadTensor2DFloat("fc3.weight");
	this->linear3Bias = &paramsLoader.LoadTensor1DFloat("fc3.bias");
}

LeNet::~LeNet()
{
	delete this->conv1Weight;
	delete this->conv1Bias;
	delete this->conv2Weight;
	delete this->conv2Bias;
	delete this->linear1Weight;
	delete this->linear1Bias;
	delete this->linear2Weight;
	delete this->linear2Bias;
	delete this->linear3Weight;
	delete this->linear3Bias;
}

int LeNet::GetResultIndex(Image& image)
{
	DL::Tensor3D<float> imageTensor = DL::Tensor3D<float>(image);
	float mean[] = { 0.5f,0.5f,0.5f };
	float std[] = { 0.5f,0.5f,0.5f };
	DL::Tensor3D<float>& input = DL::Normalize(imageTensor, mean, std);

	DL::Tensor3D<float>& conv1Result = Conv(input, 1, 0, 2, *conv1Weight, *conv1Bias);
	DL::Tensor3D<float>& conv2Result = Conv(conv1Result, 1, 0, 2, *conv2Weight, *conv2Bias);
	DL::Tensor1D<float>& flattenResult = Flatten(conv2Result);
	DL::Tensor1D<float>& linear1Result = FullConnect(flattenResult, *linear1Weight, *linear1Bias);
	DL::Tensor1D<float>& linear2Result = FullConnect(linear1Result, *linear2Weight, *linear2Bias);
	DL::Tensor1D<float>& linear3Result = Linear(linear2Result, *linear3Weight, *linear3Bias);
	DL::Tensor1D<float>& result = SoftMax(linear3Result);

	float max = -FLT_MAX;
	int index = -1;
	for (int i = 0; i < result.width; i++)
	{
		float value = result.Get(i);
		if (value > max) {
			max = value;
			index = i;
		}
	}
	delete& input;
	delete& conv1Result;
	delete& conv2Result;
	delete& flattenResult;
	delete& linear1Result;
	delete& linear2Result;
	delete& linear3Result;
	delete& result;
	return index;
}

std::string LeNet::GetResultLable(Image& image)
{
	return lables[GetResultIndex(image)];
}
